package com.infnoon.rpc.loadbalancer;

import cn.hutool.core.util.HashUtil;
import com.infnoon.rpc.model.ServiceMetaInfo;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentSkipListMap;

/**
 * 一致性 Hash 负载均衡器
 */
public class ConsistentHashLoadBalancer implements LoadBalancer {

    /**
     * 一致性 Hash 环, 存放虚拟节点
     * 线程安全优化：考虑到原代码 TreeMap 是共享的，在多线程环境下可能会导致线程安全问题，用 ConcurrentSkipListMap 替代
     */
    private final ConcurrentSkipListMap<Integer, ServiceMetaInfo> virtualNodes
            = new ConcurrentSkipListMap<>();

    /**
     * 虚拟节点数
     */
    private static final int VIRTUAL_NODE_NUM = 100;

    @Override
    public ServiceMetaInfo select(Map<String, Object> requestParams, List<ServiceMetaInfo> serviceMetaInfoList) {
        if (serviceMetaInfoList.isEmpty()) {
            return null;
        }

        // 构建虚拟节点环
        for (ServiceMetaInfo serviceMetaInfo : serviceMetaInfoList) {
            for (int i = 0; i < VIRTUAL_NODE_NUM; i ++ ) {
                int hash = getHash(serviceMetaInfo.getServiceAddress() + "#" + i);
                virtualNodes.put(hash, serviceMetaInfo);
            }
        }

        // 获取调用请求的 hash 值
        int hash = getHash(requestParams);

        // 选择最接近 且 大于等于调用请求 hash 值的虚拟节点
        Map.Entry<Integer, ServiceMetaInfo> entry = virtualNodes.ceilingEntry(hash);
        if (entry == null) {
            // 如果没有 大于等于调用请求 hash 值的虚拟节点, 返回环首部 (最小的) 的节点
            entry = virtualNodes.firstEntry();
        }
        return entry.getValue();
    }

    /**
     * Hash 算法
     * 使用 hutool 的 MurmurHash 替代默认的 HashCode
     * @param key
     * @return
     */
    private int getHash(Object key) {
        byte[] bytes = key.toString().getBytes();
        return HashUtil.murmur32(bytes);
        // return key.hashCode();
    }
}
